from langchain.llms import Tongyi
from langchain.agents import load_tools, initialize_agent, AgentType
import os
os.environ['DASHSCOPE_API_KEY'] = 'sk-965dc39b016c49ecbe29de180f4db2b6'

# 定义工具函数
def calculate_fruit_price(query):
    """解析查询并计算水果价格"""
    try:
        # 简单解析查询
        if "香蕉" in query and "苹果" in query:
            # 提取数字
            import re
            numbers = re.findall(r'\d+', query)
            if len(numbers) >= 2:
                banana_kg = int(numbers[0])
                apple_kg = int(numbers[1])
                banana_price = 8
                apple_price = 12
                total = banana_kg * banana_price + apple_kg * apple_price
                return f"香蕉{banana_kg}千克({banana_price}元/千克) + 苹果{apple_kg}千克({apple_price}元/千克) = 总共{total}元"
    except:
        pass
    return "无法计算价格，请提供明确的数量信息"

# 使用Tongyi
llm = Tongyi(
    model_name="qwen-plus",
    dashscope_api_key="sk-965dc39b016c49ecbe29de180f4db2b6"
)

# 使用内置工具
tools = load_tools(["llm-math"], llm=llm)

# 添加自定义工具
from langchain.tools import Tool
tools.append(Tool(
    name="fruit_calculator",
    func=calculate_fruit_price,
    description="计算水果价格，输入应该是包含香蕉和苹果数量的查询"
))

# 创建agent
agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)

# 执行
result = agent.run("3千克香蕉和2千克苹果需要多少钱？")
print(result)